from pathlib import Path
import json
from functools import partial
from concurrent.futures import ThreadPoolExecutor, Future
from typing import Any

import numpy as np
import cv2
from superstream import Stream


font_face = cv2.FONT_HERSHEY_COMPLEX_SMALL
font_scale = 1


def json2yolo(json_path: Path, output_path: Path, labels_info: dict) -> None:
    with json_path.open() as fp:
        data = json.load(fp)

    def normalize(data_shapes: dict, w, h):
        points = data_shapes["points"]
        x1, y1 = points[0]
        x2, y2 = points[1]
        x, y = (x1 + x2) / 2 / w, (y1 + y2) / 2 / h
        w, h = abs(x2 - x1) / w, abs(y2 - y1) / h

        label_idx = labels_info.setdefault(data_shapes["label"], len(labels_info))
        return f"{label_idx} {x:.4f} {y:.4f} {w:.4f} {h:.4f}"

    w, h = data["imageWidth"], data["imageHeight"]
    cur_normalize = partial(normalize, w=w, h=h)
    output_path.write_text(
        "\n".join(
            map(
                cur_normalize,
                filter(lambda data_dict: data_dict["shape_type"] == "rectangle", data["shapes"]),
            )
        )
    )


def main(json_folder: Path, output_folder: Path):
    output_folder.mkdir(exist_ok=True)
    labels_info = {}
    for json_file in json_folder.glob("*.json"):
        output_file = output_folder / (json_file.stem + ".txt")
        json2yolo(json_file, output_file, labels_info)

    (output_folder.resolve().parent / "label_info.txt").write_text(
        "\n".join(
            Stream(labels_info.items()).map(lambda itms: f"{itms[0]}\t=\t{itms[1]}").to_list()
        )
    )


def imwrite(path, img):
    cv2.imencode(".bmp", img)[1].tofile(path)


# 坐标转换，原始存储的是YOLOv5格式
# Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
def xywh2xyxy(yolo_content, img_w, img_h, img, output_path):
    label_idx, x, y, w, h = yolo_content
    label = label_idx
    # 边界框反归一化
    x_t = float(x) * img_w
    y_t = float(y) * img_h
    w_t = float(w) * img_w
    h_t = float(h) * img_h

    # 计算坐标
    top_left_x = x_t - w_t / 2
    top_left_y = y_t - h_t / 2
    bottom_right_x = x_t + w_t / 2
    bottom_right_y = y_t + h_t / 2

    # 绘图  rectangle()函数需要坐标为整数
    pt1 = (int(top_left_x), int(top_left_y))
    pt2 = (int(bottom_right_x), int(bottom_right_y))
    cv2.rectangle(
        img,
        pt1,
        pt2,
        (71, 127, 255),
        2,
    )
    _, base_line = cv2.getTextSize(label, font_face, font_scale, 1)
    cv2.putText(img, label, (pt1[0], pt1[1] - base_line), font_scale, font_face, (0, 0, 0), 1)

    imwrite(str(output_path), img)


def _inner_process(result: Future[tuple[Any, Path, Path, Path]]):
    img, img_path, lable_dir, output_dir = result.result()
    print(img_path)
    h, w = img.shape[:2]

    img_stem = img_path.stem
    label_path = lable_dir / (img_stem + ".txt")
    img_name = img_path.name
    output_path = output_dir / img_name

    with open(label_path) as fp:
        (
            Stream(fp.readlines())
            .map(lambda txt: txt.strip().split())
            .for_each(lambda txt: xywh2xyxy(txt, w, h, img, output_path))
        )


def imread(img_path: str, cv_flag):
    """用于解决opencv库中imread无法读取含有中文路径的imread函数"""
    return cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv_flag)


def visual_by_yolo(img_dir: Path, lable_dir: Path, output_dir: Path):
    with ThreadPoolExecutor() as exe:
        (
            Stream(img_dir.glob("*"))
            .map(
                lambda i_p: exe.submit(
                    lambda p: (imread(str(p), cv2.IMREAD_ANYCOLOR), p, lable_dir, output_dir),
                    i_p,
                )
            )
            .for_each(lambda f: f.add_done_callback(_inner_process))
        )
